package tss.dsp;

import java.util.*;
import tss.*;

class DSPHelper {
	Process process;
	Log log;

	// Used for client role
	boolean isRecording;
	Map<Integer, Collection<Message>> sentMessages, receivedMessages;

	// Used for coordinator role
	boolean isCompiling, recordingRunning;
	Collection<ProcessSnapshot> snapshots;

	public DSPHelper(Process process, Log log) {
		this.process = process;
		this.log = log;
		this.sentMessages = new Hashtable<Integer, Collection<Message>>();
		this.receivedMessages = new Hashtable<Integer, Collection<Message>>();
	}

	public void receiveMessage(Message message, int sender) {
		if (message instanceof StartMessage && sender != process.getProcessId()) {
			startRecordingInter(sender);
		} else if (message instanceof DoneMessage
				&& sender != process.getProcessId()) {
			stopRecordingInter(sender);
		} else if (message instanceof ProcessSnapshotMessage) {
			if (!recordingRunning)
				throw new IllegalStateException(
						"Snapshot received while not recording.");
			snapshots.add(((ProcessSnapshotMessage) message).getSnapshot());
		} else if (isRecording)
			for (Collection<Message> coll : receivedMessages.values())
				coll.add(message);
	}

	private void stopRecordingInter(int sender) {
		isRecording = false;
		ProcessSnapshot snapshot = new ProcessSnapshot(process.getProcessId(),
				sentMessages.get(sender), receivedMessages.get(sender), process
						.getState());
		process.sendMessage(new ProcessSnapshotMessage(snapshot), sender);
		receivedMessages.remove(sender);
		sentMessages.remove(sender);
	}

	private void startRecordingInter(int sender) {
		isRecording = true;
		sentMessages.put(sender, new ArrayList<Message>());
		receivedMessages.put(sender, new ArrayList<Message>());
	}

	public void sendMessage(Message message) {
		if (isRecording)
			for (Collection<Message> coll : sentMessages.values())
				coll.add(message);
	}

	public void startRecording() {
		if (recordingRunning)
			throw new IllegalStateException("Already recording.");

		startRecordingInter(process.getProcessId());
		recordingRunning = true;
		process.sendMessage(new StartMessage());
	}

	public Snapshot stopRecording() {
		if (isCompiling)
			throw new IllegalStateException("Already compiling snapshot");
		if (!recordingRunning)
			throw new IllegalStateException(
					"Must be recording to compile snapshot.");

		snapshots = new ArrayList<ProcessSnapshot>();
		process.sendMessage(new DoneMessage());

		stopRecordingInter(process.getProcessId());
		while (snapshots.size() < process.getGroupSize()) {
			try {
				Thread.sleep(500);
			} catch (InterruptedException _) {
			}
		}

		isCompiling = true;
		recordingRunning = false;

		return compile();
	}

	Snapshot compile() {
		Collection<Message> sentMs = new ArrayList<Message>();
		Collection<Message> recMs = new ArrayList<Message>();
		Map<Integer, Object> states = new Hashtable<Integer, Object>();

		for (ProcessSnapshot ss : snapshots) {
			sentMs.addAll(ss.getSentMessages());
			recMs.addAll(ss.getReceivedMessages());
			states.put(ss.getProcessId(), ss.getState());
		}

		sentMs.removeAll(recMs);

		return new Snapshot(states, recMs);
	}
}
